import numpy as np
import matplotlib.pyplot as plt
from scipy.special import wofz
from scipy import integrate
from tqdm import tqdm # For a helpful progress bar

# --- Plotting Setup ---
# Use LaTeX for rendering text for a professional, publication-quality look
plt.rc('text', usetex=True)
plt.rc('font', family='serif', serif=['Computer Modern'])

# Add a preamble to load necessary LaTeX math packages
plt.rcParams['text.latex.preamble'] = r'''
\usepackage{amsmath}
\usepackage{amssymb}
\usepackage{bm}
'''

# Set font sizes for labels and titles
plt.rc('xtick', labelsize=12)
plt.rc('ytick', labelsize=12)
plt.rc('axes', labelsize=14, titlesize=16)

# --- Physics & Integration Parameters ---
a = 1.0
omega = 1.0
CONST = 1.0

# Choose a few interaction times to show the trend clearly
T_omega_values = [0.5, 2.0, 10.0]

# Setup for the plot's x-axis (Ω) and the integration axis (Ω')
N_plot_points = 201
N_int_points = 401 # Use more points for the integration axis for accuracy

Omega_plot_1d = np.linspace(-3.0, 4.0, N_plot_points)
# Integration range for Ω' must be wide enough to capture the gated integrand
Omegap_int_1d = np.linspace(-5.0, 5.0, N_int_points)

# Create the 2D grid for the calculation
Omega_grid, Omegap_grid = np.meshgrid(Omega_plot_1d, Omegap_int_1d, indexing='ij')

# --- Core Calculation Function (from previous script) ---
def calculate_prob_density_2d(O_grid, Op_grid, T, a, omega, chi, chip):
    """
    Calculates the 2D probability density |M(Ω, Ω')|^2 on a grid.
    """
    eps = 1e-9
    sinh_O = np.sinh(np.pi * np.abs(O_grid) + eps)
    sinh_Op = np.sinh(np.pi * np.abs(Op_grid) + eps)

    unruh_term = (np.abs(O_grid * Op_grid) * np.exp(np.pi * (chi * O_grid + chip * Op_grid)) /
                  (sinh_O * sinh_Op))

    gate_arg = chi * O_grid + chip * Op_grid
    gaussian_gate = np.exp(-0.5 * (a * T * gate_arg)**2)

    faddeeva_arg = T * (0.5 * a * (chi * O_grid - chip * Op_grid) - omega)
    faddeeva_term = np.abs(wofz(faddeeva_arg))**2

    return CONST * unruh_term * gaussian_gate * faddeeva_term

# --- Plot Generation ---
fig, ax = plt.subplots(figsize=(10, 6))

print("Calculating Marginal Spectra...")
# Loop through interaction times
for T_omega in tqdm(T_omega_values, desc="Processing T*omega"):
    T = T_omega / omega

    # Calculate the full 2D integrand for the RR channel (χ=1, χ'=1)
    integrand_2d = calculate_prob_density_2d(Omega_grid, Omegap_grid, T, a, omega, 1, 1)

    # Calculate the marginal spectrum by integrating along the Ω' axis (axis=1)
    marginal_spectrum = integrate.simpson(integrand_2d, Omegap_int_1d, axis=1)

    # Plot the result, normalizing for easier comparison of shapes
    ax.plot(Omega_plot_1d, marginal_spectrum / np.max(marginal_spectrum),
            label=rf'$T\omega = {T_omega}$')

# --- Add Theoretical Unruh Spectrum for Comparison ---
# The single-particle spectrum for an eternal detector is proportional to the
# Bose-Einstein distribution factor multiplied by a phase-space term (Ω).
Omega_positive = np.linspace(1e-3, 4.0, 200)
# Unruh temperature T_U = a / (2*pi)
unruh_spectrum = (Omega_positive / (np.exp(2 * np.pi * Omega_positive / a) - 1))
# Normalize and plot for comparison
ax.plot(Omega_positive, unruh_spectrum / np.max(unruh_spectrum), 'k--',
        label=r'Unruh Thermal Spectrum (eternal limit)')

# --- Styling and Annotations ---
ax.set_yscale('log')

ax.set_xlabel(r'\textbf{Unruh Frequency} $\boldsymbol{\Omega}$')
ax.set_ylabel(r'\textbf{Normalized Marginal Probability} $\boldsymbol{P_R(\Omega)}$')
ax.set_title(r'\textbf{Marginal Spectrum of Right-Moving Photons (RR Channel)}')

ax.set_xlim(-3.0, 4.0)
ax.set_ylim(1e-5, 2.0)
ax.grid(True, which='both', linestyle='--', linewidth=0.5, alpha=0.7)
ax.legend(fontsize=12)

plt.tight_layout()

# --- Save the Figure ---
output_filename = 'marginal_spectra.png'
plt.savefig(output_filename, dpi=300, bbox_inches='tight')

print(f"\n✅ Plot saved successfully as '{output_filename}'")

plt.show()